use crate::{
    ast::{
        types::{BasicType as AstBasicType, MapType},
        Expression, Type,
    },
    error::{CompilerError, Result},
    llvm_backend::LLVMBackend,
};
use inkwell::{types::BasicType, values::BasicValueEnum, AddressSpace, IntPredicate};

impl<'ctx> LLVMBackend<'ctx> {
    /// Generate map.insert(key, value) method
    /// Inserts a new key-value pair or updates existing value
    pub(in crate::llvm_backend) fn generate_map_insert_method(
        &mut self,
        var_name: &str,
        arguments: &[Expression],
    ) -> Result<BasicValueEnum<'ctx>> {
        // Validate arguments
        if arguments.len() != 2 {
            return Err(CompilerError::internal(
                "insert() requires exactly 2 arguments (key, value)",
            ));
        }

        // Get the map variable pointer from variable_map
        let map_ptr = self
            .variable_map
            .get(var_name)
            .ok_or_else(|| {
                CompilerError::internal(format!("Undefined map variable: {}", var_name))
            })?
            .into_pointer_value();

        // Get the map type
        let map_type = self.variable_type_map.get(var_name).ok_or_else(|| {
            CompilerError::internal(format!("Undefined map type for variable: {}", var_name))
        })?;

        let (key_type, value_type) = if let Type::Map(MapType {
            key_type,
            value_type,
            ..
        }) = map_type
        {
            (*key_type.clone(), *value_type.clone())
        } else {
            return Err(CompilerError::internal("insert() requires a map variable"));
        };

        // Get LLVM types
        let key_llvm_type = self.nrc_type_to_llvm_type(&Some(key_type.clone()))?;
        let value_llvm_type = self.nrc_type_to_llvm_type(&Some(value_type.clone()))?;

        // Generate key and value arguments
        let new_key = self.generate_expression(&arguments[0])?;
        let new_value = self.generate_expression(&arguments[1])?;

        // Create map struct type: {keys: *K, values: *V, len: i64, cap: i64}
        let keys_ptr_type = key_llvm_type.ptr_type(AddressSpace::default());
        let values_ptr_type = value_llvm_type.ptr_type(AddressSpace::default());
        let map_struct_type = self.context.struct_type(
            &[
                keys_ptr_type.into(),
                values_ptr_type.into(),
                self.context.i64_type().into(),
                self.context.i64_type().into(),
            ],
            false,
        );

        // Load current map
        let current_map = self
            .builder
            .build_load(map_struct_type, map_ptr, "current_map")?;
        let map_struct = current_map.into_struct_value();

        // Extract fields
        let old_keys_ptr = self
            .builder
            .build_extract_value(map_struct, 0, "old_keys_ptr")?
            .into_pointer_value();
        let old_values_ptr = self
            .builder
            .build_extract_value(map_struct, 1, "old_values_ptr")?
            .into_pointer_value();
        let old_len = self
            .builder
            .build_extract_value(map_struct, 2, "old_len")?
            .into_int_value();
        let old_cap = self
            .builder
            .build_extract_value(map_struct, 3, "old_cap")?
            .into_int_value();

        // Search for existing key
        let key_index =
            self.find_key_index(old_keys_ptr, old_len, new_key, &key_type, key_llvm_type)?;

        let current_func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("No current function"))?;

        let key_found_block = self.context.append_basic_block(current_func, "key_found");
        let key_not_found_block = self
            .context
            .append_basic_block(current_func, "key_not_found");
        let merge_block = self
            .context
            .append_basic_block(current_func, "insert_merge");

        // Check if key was found (index >= 0)
        let zero = self.context.i64_type().const_int(0, false);
        let key_found =
            self.builder
                .build_int_compare(IntPredicate::SGE, key_index, zero, "key_found")?;

        self.builder
            .build_conditional_branch(key_found, key_found_block, key_not_found_block)?;

        // Key found: update existing value
        self.builder.position_at_end(key_found_block);
        let update_value_ptr = unsafe {
            self.builder.build_gep(
                value_llvm_type,
                old_values_ptr,
                &[key_index],
                "update_value_ptr",
            )?
        };
        self.builder.build_store(update_value_ptr, new_value)?;

        // Return updated map (no structural changes)
        let updated_keys_ptr = old_keys_ptr;
        let updated_values_ptr = old_values_ptr;
        let updated_len = old_len;
        let updated_cap = old_cap;

        self.builder.build_unconditional_branch(merge_block)?;

        // Key not found: insert new entry
        self.builder.position_at_end(key_not_found_block);

        // Check if we need to grow capacity
        let need_grow =
            self.builder
                .build_int_compare(IntPredicate::EQ, old_len, old_cap, "need_grow")?;

        let grow_block = self
            .context
            .append_basic_block(current_func, "grow_capacity");
        let no_grow_block = self.context.append_basic_block(current_func, "no_grow");

        self.builder
            .build_conditional_branch(need_grow, grow_block, no_grow_block)?;

        // Grow capacity block
        self.builder.position_at_end(grow_block);

        // Calculate new capacity: max(old_cap * 2, 4)
        let doubled_cap = self.builder.build_int_mul(
            old_cap,
            self.context.i64_type().const_int(2, false),
            "doubled_cap",
        )?;
        let min_cap = self.context.i64_type().const_int(4, false);
        let new_cap_grown = self
            .builder
            .build_select(
                self.builder.build_int_compare(
                    IntPredicate::UGT,
                    doubled_cap,
                    min_cap,
                    "cap_check",
                )?,
                doubled_cap,
                min_cap,
                "new_cap_grown",
            )?
            .into_int_value();

        // Allocate new keys array
        let key_size = self.get_type_size(&key_type)?;
        let new_keys_size = self.builder.build_int_mul(
            new_cap_grown,
            self.context.i64_type().const_int(key_size, false),
            "new_keys_size",
        )?;

        let malloc_func = self
            .function_map
            .get("malloc")
            .ok_or_else(|| CompilerError::internal("malloc function not found"))?;
        let new_keys_malloc =
            self.builder
                .build_call(*malloc_func, &[new_keys_size.into()], "new_keys_malloc")?;
        let new_keys_i8 = new_keys_malloc
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();
        let new_keys_ptr_grown = self
            .builder
            .build_bitcast(new_keys_i8, keys_ptr_type, "new_keys_ptr_grown")?
            .into_pointer_value();

        // Allocate new values array
        let value_size = self.get_type_size(&value_type)?;
        let new_values_size = self.builder.build_int_mul(
            new_cap_grown,
            self.context.i64_type().const_int(value_size, false),
            "new_values_size",
        )?;

        let new_values_malloc = self.builder.build_call(
            *malloc_func,
            &[new_values_size.into()],
            "new_values_malloc",
        )?;
        let new_values_i8 = new_values_malloc
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();
        let new_values_ptr_grown = self
            .builder
            .build_bitcast(new_values_i8, values_ptr_type, "new_values_ptr_grown")?
            .into_pointer_value();

        // Copy old data using memcpy (only if old_len > 0)
        let copy_needed =
            self.builder
                .build_int_compare(IntPredicate::UGT, old_len, zero, "copy_needed")?;

        let do_copy_block = self.context.append_basic_block(current_func, "do_copy");
        let skip_copy_block = self.context.append_basic_block(current_func, "skip_copy");

        self.builder
            .build_conditional_branch(copy_needed, do_copy_block, skip_copy_block)?;

        self.builder.position_at_end(do_copy_block);
        self.copy_array(old_keys_ptr, new_keys_ptr_grown, old_len, key_size)?;
        self.copy_array(old_values_ptr, new_values_ptr_grown, old_len, value_size)?;
        self.builder.build_unconditional_branch(skip_copy_block)?;

        self.builder.position_at_end(skip_copy_block);
        self.builder.build_unconditional_branch(no_grow_block)?;

        // No grow block
        self.builder.position_at_end(no_grow_block);

        let new_keys_phi = self.builder.build_phi(keys_ptr_type, "new_keys")?;
        new_keys_phi.add_incoming(&[
            (&new_keys_ptr_grown, skip_copy_block),
            (&old_keys_ptr, key_not_found_block),
        ]);
        let new_keys_ptr_insert = new_keys_phi.as_basic_value().into_pointer_value();

        let new_values_phi = self.builder.build_phi(values_ptr_type, "new_values")?;
        new_values_phi.add_incoming(&[
            (&new_values_ptr_grown, skip_copy_block),
            (&old_values_ptr, key_not_found_block),
        ]);
        let new_values_ptr_insert = new_values_phi.as_basic_value().into_pointer_value();

        let new_cap_phi = self.builder.build_phi(self.context.i64_type(), "new_cap")?;
        new_cap_phi.add_incoming(&[
            (&new_cap_grown, skip_copy_block),
            (&old_cap, key_not_found_block),
        ]);
        let new_cap_insert = new_cap_phi.as_basic_value().into_int_value();

        // Insert new key-value at position old_len
        let insert_key_ptr = unsafe {
            self.builder.build_gep(
                key_llvm_type,
                new_keys_ptr_insert,
                &[old_len],
                "insert_key_ptr",
            )?
        };
        self.builder.build_store(insert_key_ptr, new_key)?;

        let insert_value_ptr = unsafe {
            self.builder.build_gep(
                value_llvm_type,
                new_values_ptr_insert,
                &[old_len],
                "insert_value_ptr",
            )?
        };
        self.builder.build_store(insert_value_ptr, new_value)?;

        // Increment length
        let new_len_insert = self.builder.build_int_add(
            old_len,
            self.context.i64_type().const_int(1, false),
            "new_len_insert",
        )?;

        self.builder.build_unconditional_branch(merge_block)?;

        // Merge block
        self.builder.position_at_end(merge_block);

        let final_keys_phi = self.builder.build_phi(keys_ptr_type, "final_keys")?;
        final_keys_phi.add_incoming(&[
            (&updated_keys_ptr, key_found_block),
            (&new_keys_ptr_insert, no_grow_block),
        ]);
        let final_keys_ptr = final_keys_phi.as_basic_value().into_pointer_value();

        let final_values_phi = self.builder.build_phi(values_ptr_type, "final_values")?;
        final_values_phi.add_incoming(&[
            (&updated_values_ptr, key_found_block),
            (&new_values_ptr_insert, no_grow_block),
        ]);
        let final_values_ptr = final_values_phi.as_basic_value().into_pointer_value();

        let final_len_phi = self
            .builder
            .build_phi(self.context.i64_type(), "final_len")?;
        final_len_phi.add_incoming(&[
            (&updated_len, key_found_block),
            (&new_len_insert, no_grow_block),
        ]);
        let final_len = final_len_phi.as_basic_value().into_int_value();

        let final_cap_phi = self
            .builder
            .build_phi(self.context.i64_type(), "final_cap")?;
        final_cap_phi.add_incoming(&[
            (&updated_cap, key_found_block),
            (&new_cap_insert, no_grow_block),
        ]);
        let final_cap = final_cap_phi.as_basic_value().into_int_value();

        // Build final map struct
        let final_map_alloca = self.builder.build_alloca(map_struct_type, "final_map")?;

        // Store fields
        let keys_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                final_map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(0, false),
                ],
                "keys_field",
            )?
        };
        self.builder.build_store(keys_field_ptr, final_keys_ptr)?;

        let values_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                final_map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(1, false),
                ],
                "values_field",
            )?
        };
        self.builder
            .build_store(values_field_ptr, final_values_ptr)?;

        let len_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                final_map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(2, false),
                ],
                "len_field",
            )?
        };
        self.builder.build_store(len_field_ptr, final_len)?;

        let cap_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                final_map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(3, false),
                ],
                "cap_field",
            )?
        };
        self.builder.build_store(cap_field_ptr, final_cap)?;

        // Load final map and store back to variable
        let final_map =
            self.builder
                .build_load(map_struct_type, final_map_alloca, "final_map_value")?;
        self.builder.build_store(map_ptr, final_map)?;

        // Return void (represented as i32(0))
        Ok(self.context.i32_type().const_int(0, false).into())
    }

    /// Find the index of a key in the keys array
    /// Returns the index if found, or -1 if not found
    fn find_key_index(
        &mut self,
        keys_ptr: inkwell::values::PointerValue<'ctx>,
        len: inkwell::values::IntValue<'ctx>,
        search_key: BasicValueEnum<'ctx>,
        key_type: &Type,
        key_llvm_type: inkwell::types::BasicTypeEnum<'ctx>,
    ) -> Result<inkwell::values::IntValue<'ctx>> {
        let current_func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("No current function"))?;

        let loop_block = self.context.append_basic_block(current_func, "find_loop");
        let body_block = self.context.append_basic_block(current_func, "find_body");
        let found_block = self.context.append_basic_block(current_func, "find_found");
        let not_found_block = self
            .context
            .append_basic_block(current_func, "find_not_found");
        let merge_block = self.context.append_basic_block(current_func, "find_merge");

        let entry_block = self.builder.get_insert_block().unwrap();

        self.builder.build_unconditional_branch(loop_block)?;

        // Loop block
        self.builder.position_at_end(loop_block);
        let i_phi = self.builder.build_phi(self.context.i64_type(), "i")?;
        let zero = self.context.i64_type().const_int(0, false);
        i_phi.add_incoming(&[(&zero, entry_block)]);
        let i = i_phi.as_basic_value().into_int_value();

        let loop_cond = self
            .builder
            .build_int_compare(IntPredicate::SLT, i, len, "loop_cond")?;
        self.builder
            .build_conditional_branch(loop_cond, body_block, not_found_block)?;

        // Body block
        self.builder.position_at_end(body_block);
        let key_ptr = unsafe {
            self.builder
                .build_gep(key_llvm_type, keys_ptr, &[i], "key_ptr")?
        };
        let current_key = self
            .builder
            .build_load(key_llvm_type, key_ptr, "current_key")?;

        let is_equal = self.compare_keys(current_key, search_key, key_type)?;

        let i_next =
            self.builder
                .build_int_add(i, self.context.i64_type().const_int(1, false), "i_next")?;
        i_phi.add_incoming(&[(&i_next, body_block)]);

        self.builder
            .build_conditional_branch(is_equal, found_block, loop_block)?;

        // Found block
        self.builder.position_at_end(found_block);
        self.builder.build_unconditional_branch(merge_block)?;

        // Not found block
        self.builder.position_at_end(not_found_block);
        let minus_one = self.context.i64_type().const_int((-1i64) as u64, true);
        self.builder.build_unconditional_branch(merge_block)?;

        // Merge block
        self.builder.position_at_end(merge_block);
        let result_phi = self
            .builder
            .build_phi(self.context.i64_type(), "find_result")?;
        result_phi.add_incoming(&[(&i, found_block), (&minus_one, not_found_block)]);

        Ok(result_phi.as_basic_value().into_int_value())
    }

    /// Compare two keys for equality
    fn compare_keys(
        &mut self,
        key1: BasicValueEnum<'ctx>,
        key2: BasicValueEnum<'ctx>,
        key_type: &Type,
    ) -> Result<inkwell::values::IntValue<'ctx>> {
        match key_type {
            Type::Basic(basic_type) => {
                match basic_type {
                    AstBasicType::Int8
                    | AstBasicType::Int16
                    | AstBasicType::Int32
                    | AstBasicType::Int
                    | AstBasicType::Int64
                    | AstBasicType::Uint8
                    | AstBasicType::Uint16
                    | AstBasicType::Uint32
                    | AstBasicType::Uint
                    | AstBasicType::Uint64
                    | AstBasicType::Bool => {
                        let int1 = key1.into_int_value();
                        let int2 = key2.into_int_value();

                        // Handle type mismatch by casting to the wider type
                        let (cmp1, cmp2) = if int1.get_type() != int2.get_type() {
                            let width1 = int1.get_type().get_bit_width();
                            let width2 = int2.get_type().get_bit_width();

                            if width1 > width2 {
                                // Extend int2 to match int1
                                let extended = self.builder.build_int_z_extend_or_bit_cast(
                                    int2,
                                    int1.get_type(),
                                    "extend",
                                )?;
                                (int1, extended)
                            } else if width2 > width1 {
                                // Extend int1 to match int2
                                let extended = self.builder.build_int_z_extend_or_bit_cast(
                                    int1,
                                    int2.get_type(),
                                    "extend",
                                )?;
                                (extended, int2)
                            } else {
                                (int1, int2)
                            }
                        } else {
                            (int1, int2)
                        };

                        Ok(self.builder.build_int_compare(
                            IntPredicate::EQ,
                            cmp1,
                            cmp2,
                            "key_eq",
                        )?)
                    }
                    AstBasicType::String => {
                        // For strings, use strcmp
                        self.compare_strings(key1.into_pointer_value(), key2.into_pointer_value())
                    }
                    _ => Err(CompilerError::internal(
                        "Unsupported key type for comparison",
                    )),
                }
            }
            _ => Err(CompilerError::internal(
                "Only basic types supported as map keys",
            )),
        }
    }

    /// Compare two strings using strcmp
    fn compare_strings(
        &mut self,
        str1: inkwell::values::PointerValue<'ctx>,
        str2: inkwell::values::PointerValue<'ctx>,
    ) -> Result<inkwell::values::IntValue<'ctx>> {
        let strcmp_func = self
            .function_map
            .get("strcmp")
            .ok_or_else(|| CompilerError::internal("strcmp function not found"))?;

        let strcmp_result =
            self.builder
                .build_call(*strcmp_func, &[str1.into(), str2.into()], "strcmp_result")?;

        let result_int = strcmp_result
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_int_value();
        let zero = self.context.i32_type().const_int(0, false);

        Ok(self
            .builder
            .build_int_compare(IntPredicate::EQ, result_int, zero, "strings_equal")?)
    }

    /// Copy array data using memcpy
    fn copy_array(
        &mut self,
        src_ptr: inkwell::values::PointerValue<'ctx>,
        dst_ptr: inkwell::values::PointerValue<'ctx>,
        len: inkwell::values::IntValue<'ctx>,
        element_size: u64,
    ) -> Result<()> {
        let copy_size = self.builder.build_int_mul(
            len,
            self.context.i64_type().const_int(element_size, false),
            "copy_size",
        )?;

        let memcpy_intrinsic = self
            .module
            .get_function("llvm.memcpy.p0.p0.i64")
            .unwrap_or_else(|| {
                let i8_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
                let memcpy_type = self.context.void_type().fn_type(
                    &[
                        i8_ptr_type.into(),
                        i8_ptr_type.into(),
                        self.context.i64_type().into(),
                        self.context.bool_type().into(),
                    ],
                    false,
                );
                self.module
                    .add_function("llvm.memcpy.p0.p0.i64", memcpy_type, None)
            });

        let dst_i8 = self
            .builder
            .build_bitcast(
                dst_ptr,
                self.context.i8_type().ptr_type(AddressSpace::default()),
                "dst_i8",
            )?
            .into_pointer_value();

        let src_i8 = self
            .builder
            .build_bitcast(
                src_ptr,
                self.context.i8_type().ptr_type(AddressSpace::default()),
                "src_i8",
            )?
            .into_pointer_value();

        self.builder.build_call(
            memcpy_intrinsic,
            &[
                dst_i8.into(),
                src_i8.into(),
                copy_size.into(),
                self.context.bool_type().const_int(0, false).into(),
            ],
            "memcpy_call",
        )?;

        Ok(())
    }

    /// Get size of a type in bytes
    fn get_type_size(&self, ty: &Type) -> Result<u64> {
        match ty {
            Type::Basic(basic_type) => {
                Ok(match basic_type {
                    AstBasicType::Int8 | AstBasicType::Uint8 | AstBasicType::Char => 1,
                    AstBasicType::Int16 | AstBasicType::Uint16 => 2,
                    AstBasicType::Int32 | AstBasicType::Uint32 | AstBasicType::Float32 => 4,
                    AstBasicType::Int
                    | AstBasicType::Int64
                    | AstBasicType::Uint
                    | AstBasicType::Uint64
                    | AstBasicType::Float64 => 8,
                    AstBasicType::Bool => 1,
                    AstBasicType::String => 8, // Pointer
                    _ => 8,
                })
            }
            _ => Ok(8),
        }
    }
}
